from ultralytics import YOLO
import numpy as np
from typing import Optional
import queue
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from image_provider import ImageProvider
from compute_position import ComputePosition
from opt_visualizer import MergedPosePositionVisualizer

from utils import *

logger = setup_logger()

SEG_MODEL_PATH = "./merge_seg_rknn_model"
POSE_MODEL_PATH = "./merge_pose_rknn_model"

class ImageData:
    """图像数据结构"""
    def __init__(self):
        self.timestamp: float = 0.0
        self.readable_timestamp: str = ""
        self.color_img: np.ndarray = None
        self.depth_img: np.ndarray = None
        self.color_path: Optional[str] = None

class SegWorker:
    """分割工作线程"""
    def __init__(self):
        # 每个线程独立加载模型
        self.model = YOLO(SEG_MODEL_PATH, task="segment")
    
    def process(self, color_img):
        """处理分割任务"""
        try:
            seg_result = self.model(color_img)
            return self._process_result(seg_result[0])
        except Exception as e:
            print(f"Segmentation error: {e}")
            return None
    
    def _process_result(self, result):
        """快速结果处理"""
        try:
            # print(result.masks)
            if not hasattr(result, 'masks') or result.masks is None:
                return None
            masks_xy = result.masks.xy
            if len(masks_xy) == 0:
                return None
            return masks_xy[0]
        except Exception as e:
            print(f"Seg result processing error: {e}")
            return None

class PoseWorker:
    """姿态工作线程"""
    def __init__(self):
        # 每个线程独立加载模型
        self.model = YOLO(POSE_MODEL_PATH, task="pose")
    
    def process(self, color_img):
        """处理姿态任务"""
        try:
            pose_result = self.model(color_img)
            return self._process_result(pose_result[0])
        except Exception as e:
            print(f"Pose estimation error: {e}")
            return None
    
    def _process_result(self, result):
        """快速结果处理"""
        try:
            if not hasattr(result, 'keypoints') or result.keypoints is None:
                return None
            return result.keypoints
        except Exception as e:
            print(f"Pose result processing error: {e}")
            return None

class SegPose:
    def __init__(self, seg_workers=3, pose_workers=3, max_queue_size=15):
        # 首先设置关闭标志和锁 - 必须在启动线程之前
        self.shutdown_flag = False
        self.counter_lock = threading.Lock()
        
        # 创建分割工作线程池
        self.seg_executor = ThreadPoolExecutor(max_workers=seg_workers)
        self.seg_workers = [SegWorker() for _ in range(seg_workers)]
        
        # 创建姿态工作线程池
        self.pose_executor = ThreadPoolExecutor(max_workers=pose_workers)
        self.pose_workers = [PoseWorker() for _ in range(pose_workers)]
        
        # 任务管理
        self.task_queue = queue.Queue(maxsize=max_queue_size)
        self.result_queue = queue.Queue(maxsize=max_queue_size)
        self.task_counter = 0
        self.task_status = {}  # 存储任务状态 {task_id: {"seg_future": future, "pose_future": future, "image_data": image_data}}
        self.completed_tasks = {}  # 存储已完成任务 {task_id: result}
        self.next_expected_id = 0  # 下一个期待的任务ID
        
        # 启动任务处理线程 - 现在shutdown_flag已经存在
        self.processing_thread = threading.Thread(target=self._process_tasks, daemon=True)
        self.processing_thread.start()
        
        # 启动结果收集线程
        self.collection_thread = threading.Thread(target=self._collect_results, daemon=True)
        self.collection_thread.start()
    
    def put(self, image_data: ImageData):
        """添加新任务
        如果队列满，立即返回-1
        成功返回任务ID
        """
        # 生成任务ID
        with self.counter_lock:
            task_id = self.task_counter
            self.task_counter += 1
        
        # 尝试添加到任务队列（非阻塞）
        try:
            self.task_queue.put((task_id, image_data), block=False)
            return task_id
        except queue.Full:
            return -1
    
    def _process_tasks(self):
        """任务处理线程"""
        while not self.shutdown_flag:
            try:
                # 获取任务（带超时以检查关闭标志）
                task_id, image_data = self.task_queue.get(timeout=0.5)
                
                # 提交分割任务
                worker_idx = task_id % len(self.seg_workers)
                seg_future = self.seg_executor.submit(
                    self.seg_workers[worker_idx].process, 
                    image_data.color_img
                )
                
                # 提交姿态任务
                worker_idx = task_id % len(self.pose_workers)
                pose_future = self.pose_executor.submit(
                    self.pose_workers[worker_idx].process, 
                    image_data.color_img
                )
                
                # 保存任务状态
                self.task_status[task_id] = {
                    "seg_future": seg_future,
                    "pose_future": pose_future,
                    "image_data": image_data
                }
                
                # 添加回调
                seg_future.add_done_callback(lambda f, tid=task_id: self._task_completed(tid))
                pose_future.add_done_callback(lambda f, tid=task_id: self._task_completed(tid))
                
            except queue.Empty:
                continue
            except Exception as e:
                print(f"Error processing task: {e}")
    
    def _task_completed(self, task_id):
        """任务部分完成回调"""
        if task_id not in self.task_status:
            return
        
        task = self.task_status[task_id]
        
        # 检查两个任务是否都完成
        if task["seg_future"].done() and task["pose_future"].done():
            try:
                # 获取结果
                seg_result = task["seg_future"].result()
                pose_result = task["pose_future"].result()
                
                # 组装结果
                result = {
                    "seg_result": seg_result,
                    "pose_result": pose_result,
                    "image_data": task["image_data"]
                }
                
                # 保存到已完成任务
                self.completed_tasks[task_id] = result
                
                # 从任务状态中移除
                del self.task_status[task_id]
                
            except Exception as e:
                print(f"Error getting result for task {task_id}: {e}")
                # 即使失败也保存空结果
                self.completed_tasks[task_id] = {
                    "seg_result": None,
                    "pose_result": None,
                    "image_data": task["image_data"]
                }
                del self.task_status[task_id]
    
    def _collect_results(self):
        """结果收集线程（按任务ID顺序）"""
        while not self.shutdown_flag:
            # 检查是否有下一个期待的任务完成
            if self.next_expected_id in self.completed_tasks:
                # 获取结果
                result = self.completed_tasks.pop(self.next_expected_id)
                
                # 放入结果队列
                try:
                    self.result_queue.put(result, block=False)
                except queue.Full:
                    # 结果队列满时，丢弃最旧的结果
                    try:
                        self.result_queue.get(block=False)  # 丢弃一个旧结果
                    except queue.Empty:
                        pass
                    try:
                        self.result_queue.put(result, block=False)  # 放入新结果
                    except queue.Full:
                        pass  # 如果还是满，就跳过这个结果
                
                # 移动到下一个任务ID
                self.next_expected_id += 1
            else:
                # 短暂休眠避免忙等
                time.sleep(0.001)
    
    def get_result(self, block=False, timeout=None):
        """获取结果（按任务提交顺序）
        block=True: 阻塞直到有结果
        block=False: 立即返回，没有结果返回None
        """
        if block:
            try:
                return self.result_queue.get(timeout=timeout)
            except queue.Empty:
                return None
        else:
            try:
                return self.result_queue.get_nowait()
            except queue.Empty:
                return None
    
    def has_result(self):
        """检查是否有结果可获取"""
        return not self.result_queue.empty()
    
    def pending_count(self):
        """获取待处理任务数量"""
        return self.task_queue.qsize() + len(self.task_status)
    
    def shutdown(self):
        """关闭处理器"""
        self.shutdown_flag = True
        self.seg_executor.shutdown(wait=True)
        self.pose_executor.shutdown(wait=True)
        if hasattr(self, 'processing_thread'):
            self.processing_thread.join(timeout=1.0)
        if hasattr(self, 'collection_thread'):
            self.collection_thread.join(timeout=1.0)

def main():
    # 初始化SegPose处理器
    processor = SegPose(seg_workers=4, pose_workers=4, max_queue_size=20)
    
    # 初始化图像提供者
    image_provider = ImageProvider(from_folder=True)
    
    compute_position = ComputePosition()

    st_pose,st_vector=setup_standard_position(None,compute_position)

    visualizer = MergedPosePositionVisualizer(st_pose=st_pose,st_vector=st_vector)
    # 创建线程间通信队列
    result_processing_queue = queue.Queue(maxsize=20)
    
    # 线程1：图像采集和提交（10FPS）
    def image_acquisition_thread():
        """图像采集和提交线程"""
        target_interval = 0.1  # 100ms (10FPS)
        
        while True:
            start_time = time.time()
            
            # 获取图像
            img_data = image_provider.get_next_image()
            if img_data is None:
                print("Image provider finished.")
                break
            
            # 提交到处理器
            task_id = processor.put(img_data)
            if task_id == -1:
                print(f"[{time.time():.3f}] Queue full, frame dropped")
            else:
                print(f"[{time.time():.3f}] Frame submitted (ID: {task_id})")
            
            # 计算耗时并等待剩余时间
            elapsed = time.time() - start_time
            if elapsed < target_interval:
                time.sleep(target_interval - elapsed)
    
    # 线程2：结果获取（50ms间隔）
    def result_fetch_thread():
        """结果获取线程"""
        fetch_interval = 0.05  # 50ms
        
        while True:
            start_time = time.time()
            
            # 尝试获取结果
            result = processor.get_result(block=False)
            if result is not None:
                # 将结果放入处理队列
                try:
                    result_processing_queue.put(result, block=False)
                except queue.Full:
                    print("Processing queue full, result dropped")
            
            # 检查是否需要退出
            if processor.shutdown_flag and processor.pending_count() == 0:
                break
                
            # 等待剩余时间
            elapsed = time.time() - start_time
            if elapsed < fetch_interval:
                time.sleep(fetch_interval - elapsed)
    
    # 线程3：结果处理（实时处理）
    def result_processing_thread():
        """结果处理线程"""
        while True:
            try:
                # 阻塞获取结果（带超时以检查退出条件）
                result = result_processing_queue.get(timeout=0.5)
                pose = None
                vector = None
                if result['image_data'] is not None and result['seg_result'] is not None and result['pose_result'] is not None:
                    time_stamp=result['image_data'].timestamp
                    color_img = result['image_data'].color_img
                    depth_img = result['image_data'].depth_img
                    seg_result=result['seg_result']
                    pose_result=result['pose_result']
                    pose,vector = compute_position.process_image_pair(time_stamp,color_img,depth_img,seg_result,pose_result)
                print(pose)
                print(vector)
                # 标记任务完成
                result_processing_queue.task_done()
                
            except queue.Empty:
                # 检查是否需要退出
                if processor.shutdown_flag and processor.pending_count() == 0:
                    break
    
    # 创建并启动线程
    acquisition_thread = threading.Thread(target=image_acquisition_thread, name="ImageAcquisition")
    fetch_thread = threading.Thread(target=result_fetch_thread, name="ResultFetch")
    processing_thread = threading.Thread(target=result_processing_thread, name="ResultProcessing")
    
    acquisition_thread.start()
    fetch_thread.start()
    processing_thread.start()
    
    # 等待图像采集线程结束
    acquisition_thread.join()
    print("Image acquisition finished. Waiting for processing to complete...")
    
    # 等待剩余任务处理完成
    processor.shutdown_flag = True
    
    # 等待结果获取线程结束
    fetch_thread.join()
    print("Result fetching finished.")
    
    # 等待结果处理线程结束
    processing_thread.join()
    print("Result processing finished.")
    
    # 正式关闭处理器
    processor.shutdown()
    print("Processor shutdown complete.")

if __name__ == "__main__":
    main()